# -*- coding:utf-8 -*-
# @Time: 2022/9/21 8:48
# @Author: ChenmingSong
# @Email: SongCM@CATL.com
# @File: panda2coco.py
import cv2
import json


def read_json_vehicle(json_file):
    vehicle_cats = ['motorcycle', 'small car', 'bicycle', 'unsure', 'baby carriage', 'midsize car', 'large car',
                    'electric car', 'tricycle']
    vehicle_json = {"images": [],
                    "categories": [{
                        "supercategory": "vehicle",
                        "id": 1,
                        "name": "vehicle"
                    }],
                    "annotations": []
                    }

    with open(json_file, "r", encoding="utf-8") as f:
        data = json.load(f)
        annotations_idx = 1
        for key, value in data.items():
            images_file_name = key
            images_id = value["image id"]
            images_height = value['image size']['height']
            images_width = value['image size']['width']
            vehicle_json["images"].append({
                "height": images_height,
                "width": images_width,
                "id": images_id,
                "file_name": images_file_name
            })
            for object in value["objects list"]:
                print(object)
                cat = object["category"]
                if cat in vehicle_cats:  # 如果类别在所需要的类别中，执行下面的操作，提取对应bbox
                    tl = object['rect']['tl']
                    br = object['rect']['br']
                    tl_x = tl['x'] * images_width
                    tl_y = tl['y'] * images_height
                    # br_x = br['x']
                    # br_y = br['y']
                    rect_w = (br['x'] - tl['x']) * images_width
                    rect_h = (br['y'] - tl['y']) * images_height
                    rect_x = int(tl_x + 0.5)
                    rect_y = int(tl_y + 0.5)
                    vehicle_json["annotations"].append(
                        {"iscrowd": 0,
                         "image_id": images_id,
                         "bbox": [
                             float(rect_x),
                             float(rect_y),
                             float(int(rect_w+0.5)),
                             float(int(rect_h+0.5))
                         ],
                         # "category_id": vehicle_cats.index(cat) + 1,
                         "category_id": 1,
                         "id": annotations_idx,
                         "area": int(rect_h * rect_w + 0.5)
                         })
                    annotations_idx = annotations_idx + 1

    with open("gigavision_vehicle_coco_format.json", "w", encoding="utf-8") as save_f:
        json.dump(vehicle_json, save_f, indent=4, ensure_ascii=False)

# coco数据格式
def read_json_all(json_file_vehicle, json_file_person):
    vehicle_cats = ['motorcycle', 'small car', 'bicycle', 'unsure', 'baby carriage', 'midsize car', 'large car',
                    'electric car', 'tricycle']
    person_cats = ['person']
    vehicle_json = {"images": [],
                    "categories": [{
                        "supercategory": "vehicle",
                        "id": 1,
                        "name": "vehicle"
                    }, {
                        "supercategory": "person",
                        "id": 2,
                        "name": "person"
                    }],
                    "annotations": []
                    }

    with open(json_file_vehicle, "r", encoding="utf-8") as f:
        data = json.load(f)
        annotations_idx = 1
        for key, value in data.items():
            images_file_name = key
            images_id = value["image id"]
            images_height = value['image size']['height']
            images_width = value['image size']['width']
            vehicle_json["images"].append({
                "height": images_height,
                "width": images_width,
                "id": images_id,
                "file_name": images_file_name
            })
            for object in value["objects list"]:
                print(object)
                cat = object["category"]
                if cat in vehicle_cats:  # 如果类别在所需要的类别中，执行下面的操作，提取对应bbox
                    tl = object['rect']['tl']
                    br = object['rect']['br']
                    tl_x = tl['x'] * images_width
                    tl_y = tl['y'] * images_height
                    # br_x = br['x']
                    # br_y = br['y']
                    rect_w = (br['x'] - tl['x']) * images_width
                    rect_h = (br['y'] - tl['y']) * images_height
                    rect_x = int(tl_x + 0.5)
                    rect_y = int(tl_y + 0.5)
                    vehicle_json["annotations"].append(
                        {"iscrowd": 0,
                         "image_id": images_id,
                         "bbox": [
                             float(rect_x),
                             float(rect_y),
                             float(int(rect_w+0.5)),
                             float(int(rect_h+0.5))
                         ],
                         # "category_id": vehicle_cats.index(cat) + 1,
                         "category_id": 1,
                         "id": annotations_idx,
                         "area": int(rect_h * rect_w + 0.5)
                         })
                    annotations_idx = annotations_idx + 1

    with open(json_file_person, "r", encoding="utf-8") as f:
        data = json.load(f)
        annotations_idx = 1
        for key, value in data.items():
            images_file_name = key
            images_id = value["image id"]
            images_height = value['image size']['height']
            images_width = value['image size']['width']
            # vehicle_json["images"].append({
            #     "height": images_height,
            #     "width": images_width,
            #     "id": images_id,
            #     "file_name": images_file_name
            # })
            for object in value["objects list"]:
                print(object)
                cat = object["category"]
                if cat in person_cats:  # 如果类别在所需要的类别中，执行下面的操作，提取对应bbox
                    tl = object['rects']['full body']['tl']
                    br = object['rects']['full body']['br']
                    tl_x = tl['x'] * images_width
                    tl_y = tl['y'] * images_height
                    # br_x = br['x']
                    # br_y = br['y']
                    rect_w = (br['x'] - tl['x']) * images_width
                    rect_h = (br['y'] - tl['y']) * images_height
                    rect_x = int(tl_x + 0.5)
                    rect_y = int(tl_y + 0.5)
                    vehicle_json["annotations"].append(
                        {"iscrowd": 0,
                         "image_id": images_id,
                         "bbox": [
                             float(rect_x),
                             float(rect_y),
                             float(int(rect_w+0.5)),
                             float(int(rect_h+0.5))
                         ],
                         # "category_id": vehicle_cats.index(cat) + 1,
                         "category_id": 1,
                         "id": annotations_idx,
                         "area": int(rect_h * rect_w + 0.5)
                         })
                    annotations_idx = annotations_idx + 1

    with open("gigavision_all_coco_format.json", "w", encoding="utf-8") as save_f:
        json.dump(vehicle_json, save_f, indent=4, ensure_ascii=False)


if __name__ == '__main__':
    # 训练集一共是390张图片。
    read_json_all(json_file_vehicle="vehicle_bbox_train.json", json_file_person="person_bbox_train.json")
############################################## wbf
# -*- coding:utf-8 -*-
# @Time: 2022/9/22 14:26
# @Author: ChenmingSong
# @Email: SongCM@CATL.com
# @File: wbf.py
from ensemble_boxes import *

boxes_list = [[
    [0.00, 0.51, 0.81, 0.91],
    [0.10, 0.31, 0.71, 0.61],
    [0.01, 0.32, 0.83, 0.93],
    [0.02, 0.53, 0.11, 0.94],
    [0.03, 0.24, 0.12, 0.35],
],[
    [0.04, 0.56, 0.84, 0.92],
    [0.12, 0.33, 0.72, 0.64],
    [0.38, 0.66, 0.79, 0.95],
    [0.08, 0.49, 0.21, 0.89],
]]
scores_list = [[0.9, 0.8, 0.2, 0.4, 0.7], [0.5, 0.8, 0.7, 0.3]]
labels_list = [[0, 1, 0, 1, 1], [1, 1, 1, 0]]
weights = [2, 1]

iou_thr = 0.5
skip_box_thr = 0.0001
sigma = 0.1

boxes, scores, labels = nms(boxes_list, scores_list, labels_list, weights=weights, iou_thr=iou_thr)
boxes, scores, labels = soft_nms(boxes_list, scores_list, labels_list, weights=weights, iou_thr=iou_thr, sigma=sigma, thresh=skip_box_thr)
boxes, scores, labels = non_maximum_weighted(boxes_list, scores_list, labels_list, weights=weights, iou_thr=iou_thr, skip_box_thr=skip_box_thr)
boxes, scores, labels = weighted_boxes_fusion(boxes_list, scores_list, labels_list, weights=weights, iou_thr=iou_thr, skip_box_thr=skip_box_thr)
print(boxes)
print(scores)
print(labels)

################################ move data
# -*- coding:utf-8 -*-
# @Time: 2022/9/22 13:48
# @Author: ChenmingSong
# @Email: SongCM@CATL.com
# @File: move_data.py
import os
import cv2
import os.path as osp
import shutil


def move_images(label_folder, image_src_folder, image_save_folder):
    label_names = os.listdir(label_folder)
    for label_name in label_names:
        image_src_path = osp.join(image_src_folder, label_name.split(".")[0] + ".jpg")
        shutil.copy2(image_src_path, image_save_folder)


# if __name__ == '__main__':
#     move_images(r"E:\ddddemo\sahi-main\runs\coco2yolov5\exp\train", )